
from collections import defaultdict
from copy import deepcopy
import random,math


SIZE = int(input("选择3x3,5x5：3 or 5")) # 网格大小
# TRY_NUMBER = int(input("难度：1000 — 5000："))
TRY_NUMBER = 4000 # 难度调节



class CONSTANTS:
    X = 1
    O = -1
    # GAME result
    DRAW = 0 #平局
    CONTINUE = 255 # 继续游戏

    # 3 x 3的三连数组
    wincombinations = [
        [0,1,2],# row
        [3,4,5],
        [6,7,8],
        [0,3,6], # col
        [1,4,7],
        [2,5,8],
        [0,4,8],# dig
        [2,4,6] 
    ]


    wincombinations_5x5 = [
        [0,1,2,3],#row
        [1,2,3,4],
        [5,6,7,8],
        [6,7,8,9],
        [10,11,12,13],
        [11,12,13,14],
        [15,16,17,18],
        [16,17,18,19],
        [20,21,22,23],
        [21,22,23,24],
        [0,5,10,15],#col
        [5,10,15,20],
        [1,6,11,16],
        [6,11,16,21],
        [2,7,12,17],
        [7,12,17,22],
        [3,8,13,18],
        [8,13,18,23],
        [4,9,14,19],
        [9,14,19,24],
        [1,7,13,19], #dig
        [5,11,17,23],
        [0,6,12,18],
        [6,12,18,24],
        [3,7,11,15],
        [9,13,17,21],
        [4,8,12,16],
        [8,12,16,20]

    ]

class ChessMove:
    def __init__(self,pos,value):
        self.pos = pos
        self.value = value
    def __repr__(self):
        return "{},{}".format(self.pos,self.value)


class ChessState:

    def __init__(self,board,to_move):
        self.board = board
        self.to_move = to_move

    @property
    def game_result(self):
        # 返回胜者

        result = CONSTANTS.CONTINUE
        if SIZE == 3:
            for a,b,c in CONSTANTS.wincombinations:
                if  self.board[a] == self.board[b] == self.board[c] ==  CONSTANTS.X:
                    result =  CONSTANTS.X
                    break
                elif self.board[a] == self.board[b] == self.board[c] ==  CONSTANTS.O:
                    result =  CONSTANTS.O
                    break
        elif SIZE == 5:
            for a,b,c,d in CONSTANTS.wincombinations_5x5:
                if  self.board[a] == self.board[b] == self.board[c] == self.board[d]== CONSTANTS.X:
                    result =  CONSTANTS.X
                    break
                elif self.board[a] == self.board[b] == self.board[c] == self.board[d] ==  CONSTANTS.O:
                    result =  CONSTANTS.O
                    break
       
        
        if result == CONSTANTS.CONTINUE and self.get_legal_actions() == []:
            result = CONSTANTS.DRAW # 平局

        return result


    def is_gameover(self):
        return self.game_result != CONSTANTS.CONTINUE

    def is_move_legal(self,move):
        if self.to_move != move.value :
            return False
        if not (0 <=  move.pos < SIZE*SIZE):
            return False
        if self.board[move.pos] != 0:
            return False

        return True


    def get_legal_actions(self):
        
        valid_moves = []
        
        for i in range(SIZE * SIZE):
            if self.board[i] == 0:
                move = ChessMove(i,self.to_move)

                valid_moves.append(move)

        return valid_moves


    def chessmove(self,move):
        if not self.is_move_legal(move):
            raise ValueError("ERROR")
        new_board = deepcopy(self.board)
        new_board[move.pos] = move.value

        to_move = CONSTANTS.O if self.to_move == CONSTANTS.X else CONSTANTS.X

        return ChessState(new_board,to_move)


class MCTSNode:

    def __init__(self,state,parent = None):
        self.state = state
        self.parent = parent
        self.children = []
        self._results = defaultdict(int)
        self._N = 0
        self.untried_actions = self.state.get_legal_actions()


    # @property
    # def untried_actions(self):
    #     if not hasattr(self,'_untried_actions'):
    #         self._untried_actions = self.state.get_legal_actions()
    #     return self._untried_actions

    @property
    def Q(self):
        #收益
        # 平局没管
        wins = self._results[self.parent.state.to_move] #父节点的tomove就是本节点
        loses = self._results[- self.parent.state.to_move]
        
        return wins - loses

    @property
    def N(self):
        return self._N

    def is_terminal_node(self):
        return self.state.is_gameover()

    def is_full_expanded(self):
        return len(self.untried_actions) == 0

    def expand(self):
        action = self.untried_actions.pop()
        t_state = self.state.chessmove(action)
        t_node = MCTSNode(t_state,self)#当前局面先走一步
        self.children.append(t_node)
        return t_node

    def rollout(self):
        #走一步之后随机走子模拟到结束
        cur_state = self.state
        while not cur_state.is_gameover():
            possible_moves = cur_state.get_legal_actions()
            action = random.choice(possible_moves)
            cur_state = cur_state.chessmove(action)

        return cur_state.game_result

    def back(self,reward):
        self._N += 1
        self._results[reward] += 1 # 简单的把每种结局统计次数
        if self.parent:
            self.parent.back(reward)

    def best_child(self,c_param = 1 / math.sqrt(2)):
        choices_weight = [(child.Q / child.N + c_param * math.sqrt(2 * math.log(self.N) / child.N))
         for child in self.children] # UCB
        max_choice = max(choices_weight)
        return self.children[choices_weight.index(max_choice)]


class MCTS:
    def __init__(self,node):
        self.node = node

    def tree_policy(self):
        #选择、扩展
        cur_node = self.node
        while not cur_node.is_terminal_node():
            if not cur_node.is_full_expanded():
                return cur_node.expand()
            else:
                cur_node = cur_node.best_child()#综合考量
        return cur_node

    def best_action(self,s_number):
        for i in range(s_number):
            try_node = self.tree_policy()
            reward = try_node.rollout() # 实际返回的是结局胜、负的对象
            try_node.back(reward)

        return self.node.best_child(c_param = 0)#这里只选平均收益最大的




def print_3x3(board):
    print(board[0] ,end=" ")
    print(board[1] ,end=" ")
    print(board[2])

    print(board[3] ,end=" ")
    print(board[4] ,end=" ")
    print(board[5])

    print(board[6] ,end=" ")
    print(board[7] ,end=" ")
    print(board[8])


def print_5x5(board):
    print(board[0] ,end=" ")
    print(board[1] ,end=" ")
    print(board[2] ,end=" ")
    print(board[3] ,end=" ")
    print(board[4])

    print(board[5] ,end=" ")
    print(board[6] ,end=" ")
    print(board[7] ,end=" ")
    print(board[8] ,end=" ")
    print(board[9])

    print(board[10] ,end=" ")
    print(board[11] ,end=" ")
    print(board[12] ,end=" ")
    print(board[13] ,end=" ")
    print(board[14])

    print(board[15] ,end=" ")
    print(board[16] ,end=" ")
    print(board[17] ,end=" ")
    print(board[18] ,end=" ")
    print(board[19])

    print(board[20] ,end=" ")
    print(board[21] ,end=" ")
    print(board[22] ,end=" ")
    print(board[23] ,end=" ")
    print(board[24])





def graphics(grid):

    board = ['' for _ in range(SIZE * SIZE)]
    for i in range(SIZE * SIZE):
        if grid[i] == 1:
            board[i] = "X"
        elif grid[i] == -1:
            board[i] = "O"
        else:
            board[i] = "-"

    if SIZE == 3:
        print_3x3(board)
    elif SIZE == 5:
        print_5x5(board)    
    

    print("------------------------")
        




def get_action(state):
    try:
        location = input("输入行和列（下标0开始） ，使用空格分隔：")
        r ,c =  location.split(" ")
        pos = int(r) * SIZE + int(c)
        move = ChessMove(pos,CONSTANTS.O)
    except Exception as e:
        move = -1
    if move == -1 or not state.is_move_legal(move):
        print("invalid move")
        move = get_action(state)
    return move


def judge(state):
    if state.is_gameover():
        if state.game_result ==CONSTANTS.X:
            print("You lose!")
        elif state.game_result == CONSTANTS.O:
            print("You Win!")
        elif state.game_result == CONSTANTS.DRAW:
            print("draw!")
    else:
        return CONSTANTS.CONTINUE







# 初始化定了电脑X先走
def init():
    init_board = [0 for _ in range(SIZE * SIZE)]
    init_board_state = ChessState(init_board,CONSTANTS.X)#x先走
    root = MCTSNode(init_board_state,None)
    mcts = MCTS(root)
    best_node = mcts.best_action(TRY_NUMBER)
    c_state = best_node.state
    c_board = c_state.board
    return c_state,c_board



c_state,c_board = init()
graphics(c_board)

while True:
    move = get_action(c_state)
    c_state = c_state.chessmove(move) # 玩家下棋
    c_board = c_state.board
    graphics(c_board)
    
    if judge(c_state)!= CONSTANTS.CONTINUE:
        break

    board_state = ChessState(c_board,CONSTANTS.X) # 电脑走
    node = MCTSNode(board_state,None)
    m = MCTS(node)
    best_node = m.best_action(TRY_NUMBER)
    c_state = best_node.state
    c_board = c_state.board
    graphics(c_board)
    if judge(c_state)!= CONSTANTS.CONTINUE:
        break






    

